[ExecuTorch][Vulkan] Add et_vk.fused_ce kernels (fused cross-entropy loss + dlogits)#20941
[ExecuTorch][Vulkan] Add et_vk.fused_ce kernels (fused cross-entropy loss + dlogits)#20941JCNTH wants to merge 3 commits into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20941
Note: Links to docs will display an error until the docs builds have been completed. ❌ 9 New Failures, 1 Unrelated FailureAs of commit d774a33 with merge base 7013c8d ( NEW FAILURES - The following jobs have failed:
BROKEN TRUNK - The following job failed but were present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
This PR needs a
|
Stack from ghstack (oldest at bottom):
Vulkan GLSL kernels + handler for the training custom op
et_vk.fused_ce(fused CE: per-row online-softmax loss + dlogits, then an N->1 loss sum). Per-row shared-memory reduction structured likeglsl/reduce_per_row_buffer.glsl(NWORKERS=64, one workgroup/row, barrier tree-combine); two dispatch nodes shareloss_partial, ordered by the runtime per-vTensorbarrier (same pattern asimpl/SDPA.cppQK->softmax->AV). AOT registration already exists under the shared Vulkan partitioner.@exported-using-ghexport
Differential Revision: D111761780
Differential Revision: D111761780